// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#pragma once

#include <pollux/expression/expr.h>
#include <pollux/expression/register_special_form.h>
#include <pollux/parse/expressions.h>
#include <pollux/parse/expressions_parser.h>
#include <pollux/parse/type_resolver.h>
#include <pollux/vector/vector_maker.h>

namespace kumo::pollux::functions::test {

class FunctionBenchmarkBase {
 public:
  FunctionBenchmarkBase() {
    parse::registerTypeResolver();
    exec::registerFunctionCallToSpecialForms();
  }

  void setTimezone(const std::string& value) {
    queryCtx_->testingOverrideConfigUnsafe({
        {core::QueryConfig::kSessionTimezone, value},
    });
  }

  void setAdjustTimestampToTimezone(const std::string& value) {
    queryCtx_->testingOverrideConfigUnsafe(
        {{core::QueryConfig::kAdjustTimestampToTimezone, value}});
  }

  exec::ExprSet compileExpression(
      const std::string& text,
      const TypePtr& rowType) {
    auto untyped = parse::parseExpr(text, options_);
    auto typed =
        core::Expressions::inferTypes(untyped, rowType, execCtx_.pool());
    return exec::ExprSet({typed}, &execCtx_);
  }

  void evaluate(
      exec::ExprSet& exprSet,
      const RowVectorPtr& data,
      const SelectivityVector& rows,
      VectorPtr& result) {
    exec::EvalCtx evalCtx(&execCtx_, &exprSet, data.get());
    std::vector<VectorPtr> results{result};
    exprSet.eval(rows, evalCtx, results);

    // If result was nullptr, we need to pick up the value that was created.
    if (!result) {
      result = results[0];
    }
  }

  VectorPtr evaluate(
      exec::ExprSet& exprSet,
      const RowVectorPtr& data,
      const SelectivityVector& rows) {
    VectorPtr result;
    evaluate(exprSet, data, rows, result);
    return result;
  }

  VectorPtr evaluate(exec::ExprSet& exprSet, const RowVectorPtr& data) {
    SelectivityVector rows(data->size());
    return evaluate(exprSet, data, rows);
  }

  VectorPtr evaluate(const std::string& expression, const RowVectorPtr& data) {
    auto exprSet = compileExpression(expression, as_row_type(data->type()));
    return evaluate(exprSet, data);
  }

  kumo::pollux::VectorMaker& maker() {
    return vectorMaker_;
  }

  memory::MemoryPool* pool() {
    return execCtx_.pool();
  }

 protected:
  std::shared_ptr<core::QueryCtx> queryCtx_{core::QueryCtx::create()};
  std::shared_ptr<memory::MemoryPool> pool_{
      memory::memoryManager()->addLeafPool()};
  core::ExecCtx execCtx_{pool_.get(), queryCtx_.get()};
  kumo::pollux::VectorMaker vectorMaker_{execCtx_.pool()};
  parse::ParseOptions options_;
};
} // namespace kumo::pollux::functions::test
